import onnxruntime as ort
import cv2 as cv
import numpy as np

# 打印设备
device_name = ort.get_device()
print("Device:", device_name)

# 创建 ONNX 推理会话
session = ort.InferenceSession(
    "../../game/model/best.onnx",
    providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
)

# 读取图像并转换格式
image = cv.imread("2.png")
image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
image = cv.resize(image, (640, 640))  # 按照模型要求调整尺寸

# 预处理：归一化 + 变形
image = image.astype(np.float32) / 255.0
image = np.transpose(image, (2, 0, 1))  # HWC -> CHW
image = np.expand_dims(image, axis=0)  # CHW -> NCHW

# 推理
inputs = {session.get_inputs()[0].name: image}
outputs = session.run(None, inputs)

print(outputs)
